package logger

import (
	"context"
	"github.com/sirupsen/logrus"
	"gorm.io/gorm/logger"
	"gorm.io/gorm/utils"
	"gqlgen-test2/server/constant"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"time"
)

var gormSourceDir string

func init() {
	_, file, _, _ := runtime.Caller(0)
	gormSourceDir = regexp.MustCompile(`utils.utils\.go`).ReplaceAllString(file, "")
}

type GormLogger struct {
	LogI          *logrus.Entry
	LogLevel      logger.LogLevel
	SlowThreshold time.Duration
}

func (g GormLogger) LogMode(level logger.LogLevel) logger.Interface {
	g.LogLevel = level
	return g
}

func (g GormLogger) Info(ctx context.Context, s string, i ...interface{}) {
	var currentLog *logrus.Entry
	l, ok := ctx.Value(constant.CONTEXT_LOG).(*logrus.Entry)
	if ok {
		currentLog = l
	} else {
		currentLog = g.LogI
	}
	if g.LogLevel >= logger.Info {
		currentLog.Info(s, append([]interface{}{utils.FileWithLineNum()}, i...))
	}
}

func (g GormLogger) Warn(ctx context.Context, s string, i ...interface{}) {
	var currentLog *logrus.Entry
	l, ok := ctx.Value(constant.CONTEXT_LOG).(*logrus.Entry)
	if ok {
		currentLog = l
	} else {
		currentLog = g.LogI
	}
	if g.LogLevel >= logger.Warn {
		currentLog.Warn(s, append([]interface{}{utils.FileWithLineNum()}, i...))
	}
}

func (g GormLogger) Error(ctx context.Context, s string, i ...interface{}) {
	var currentLog *logrus.Entry
	l, ok := ctx.Value(constant.CONTEXT_LOG).(*logrus.Entry)
	if ok {
		currentLog = l
	} else {
		currentLog = g.LogI
	}
	if g.LogLevel >= logger.Error {
		currentLog.Error(s, append([]interface{}{utils.FileWithLineNum()}, i...))
	}
}

func (g GormLogger) Trace(ctx context.Context, begin time.Time, fc func() (string, int64), err error) {
	var currentLog *logrus.Entry
	l, ok := ctx.Value(constant.CONTEXT_LOG).(*logrus.Entry)
	if ok {
		currentLog = l
	} else {
		currentLog = g.LogI
	}
	if g.LogLevel > 0 {
		elapsed := time.Since(begin)
		switch {
		case err != nil && g.LogLevel >= logger.Error:
			sql, rows := fc()
			if rows == -1 {
				currentLog.WithFields(logrus.Fields{
					"timeUse":  elapsed.String(),
					"location": utils.FileWithLineNum(),
					"err":      err,
					"rows":     "-",
				}).Debug(sql)
			} else {
				currentLog.WithFields(logrus.Fields{
					"timeUse":  elapsed.String(),
					"location": utils.FileWithLineNum(),
					"err":      err,
					"rows":     rows,
				}).Debug(sql)
			}
		case elapsed > g.SlowThreshold && g.SlowThreshold != 0 && g.LogLevel >= logger.Warn:
			sql, rows := fc()
			if rows == -1 {
				currentLog.WithFields(logrus.Fields{
					"timeUse":  elapsed.String(),
					"location": utils.FileWithLineNum(),
					"rows":     "-",
				}).Debug(sql)
			} else {
				currentLog.WithFields(logrus.Fields{
					"timeUse":  elapsed.String(),
					"location": utils.FileWithLineNum(),
					"rows":     rows,
				}).Debug(sql)
			}
		case g.LogLevel >= logger.Info:
			sql, rows := fc()
			if rows == -1 {
				currentLog.WithFields(logrus.Fields{
					"timeUse":  elapsed.String(),
					"location": utils.FileWithLineNum(),
					"rows":     "-",
				}).Debug(sql)
			} else {
				currentLog.WithFields(logrus.Fields{
					"timeUse":  elapsed.String(),
					"location": utils.FileWithLineNum(),
					"rows":     rows,
				}).Debug(sql)
			}
		}
	}
}

func FileWithLineNum() string {
	for i := 2; i < 15; i++ {
		_, file, line, ok := runtime.Caller(i)

		if ok && (!strings.HasPrefix(file, gormSourceDir) || strings.HasSuffix(file, "_test.go")) {
			return strconv.FormatInt(int64(line), 10)
		}
	}
	return ""
}
